import torch
from torch import nn, optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader,Subset
import torchvision.models as models
from tqdm.notebook import tqdm
from torch.utils.data import Dataset
import numpy as np
import itertools



def freeze_layers(model, freeze_prefixes):
    """
    冻结模型中指定前缀的层。
    
    参数:
    - model: 要冻结层的模型。
    - freeze_prefixes: 一个列表，包含要冻结的层的前缀。
    """


    for name, param in model.named_parameters():
        freeze = False
        for prefix in freeze_prefixes:
            if name.startswith(prefix):
                param.requires_grad = False  # 冻结参数，不进行梯度更新
                freeze = True
                print(f"----Freezing {name}")
        if not freeze:
            print(name)  # 打印未冻结的层名称
            param.requires_grad = True


#---------------------损失函数相关

def smooth_l1_loss(x, beta=1.0):
    """
    计算 Smooth L1 Loss（Huber Loss）。

    参数:
    x -- 预测值与真实值之间的差异，数值或numpy数组
    beta -- 转换点，当误差小于 beta 时使用平方误差，大于 beta 时使用绝对误差，默认值为 1.0

    返回:
    Smooth L1 Loss 的计算结果
    """
    # 计算绝对值
    abs_x = np.abs(x)
    # 计算 Smooth L1 Loss
    loss = np.where(abs_x < beta, 0.5 * x**2, beta * (abs_x - 0.5 * beta))
    return loss


'''def get_class_frequency(dataloader):
    """
    从DataLoader中获取类别及其频率的字典。
    
    参数:
    dataloader -- DataLoader实例。
    
    返回:
    class_frequency -- 类别及其频率的字典。
    """
    class_frequency = {}
    total_samples = 0  # 总样本数
    for inputs, labels in dataloader:
        labels = labels.cpu().numpy()  # 将标签转换为numpy数组
        for label in labels:
            if label not in class_frequency:
                class_frequency[label] = 1
            else:
                class_frequency[label] += 1
            total_samples += 1  # 每次遇到一个标签，总样本数加1
    
    # 计算频率
    for label in class_frequency:
        class_frequency[label] = class_frequency[label] / total_samples
    
    return class_frequency'''

